import math
import numpy
class Nodes:
    """
    Types must be:
    1 for tough node (zhestkaya zadelka)
    2 for hinge node (sharnir)


    """
    def __init__(self, NodeNumber, NodeType, NodeCoor):
        self.NodeNumber 	= NodeNumber
        self.NodeType 		= NodeType
        self.NodeCoor		= NodeCoor

class Beams:

    def __init__(self, BeamNumber, StartNode, EndNode, ElasticModulus, Area, InertiaMoment):
        self.BeamNumber 	= BeamNumber
        self.StartNode		= StartNode
        self.EndNode		= EndNode
        self.ElasticModulus	= ElasticModulus
        self.Area			= Area
        self.InertiaMoment	= InertiaMoment
        self.BeamType = self.StartNode.NodeType + self.EndNode.NodeType

        def CalcBeamLength(StartNode, EndNode):
            XN=StartNode.NodeCoor	[0]
            XK=EndNode.NodeCoor		[0]
            YN=StartNode.NodeCoor	[1]
            YK=EndNode.NodeCoor		[1]
            Length = math.sqrt((XK - XN)**2 + (YK - YN)**2)
            return Length

        self.BeamLength		= CalcBeamLength(StartNode, EndNode)



def CreateLocalStiffnessMatrix(Beam):
    """CreateStiffnessMatrix"""
    EFL		=Beam.ElasticModulus * Beam.Area / Beam.BeamLength
    EFL2	=Beam.ElasticModulus * Beam.Area / (Beam.BeamLength**2)
    EFL3	=Beam.ElasticModulus * Beam.Area / (Beam.BeamLength**3)
    SMatrix =numpy.array()
    if Beam.BeamType 	== 2:
        SMatrix = numpy.array ([(EFL,	    0, 		0,     -EFL,      0,        0  ),
                            (0, 	 12*EFL3, 6*EFL2,    0,   -12*EFL3,   6*EFL2),
                            (0,     6*EFL2,  4*EFL, 	 0,	  -6*EFL2,    2*EFL),
                            (-EFL, 	0,   	0,	    EFL, 	  0,		  0  ),
                            (0,	 -12*EFL3, -6*EFL2,  0,  	12*EFL3,  -6*EFL2),
                            (0,	  6*EFL2,  2*EFL,    0,		-6*EFL2,   4*EFL)
                            ])

    elif Beam.BeamType 	== 3:
        SMatrix = numpy.array ([(EFL,      0,      0,   -EFL,    0   ),
                             ( 0,    3*EFL3, 3*EFL2,   0,  -3*EFL3),
                             ( 0,    3*EFL2, 3*EFL,    0,  -3*EFL2),
                             (-EFL,     0,      0,    EFL,    0   ),
                             ( 0,   -3*EFL3,-3*EFL2,   0,   3*EFL3),
                            ])

    elif Beam.BeamType	== 4:
        SMatrix = numpy.array ([(EFL,      0,    -EFL,    0),
                             ( 0,       0,      0,     0),
                             (-EFL,     0,     EFL,    0),
                             ( 0,       0,      0,     0)
                            ])

    return SMatrix

def CreateGlobalStiffnessMatrix(Beam):
    """Convert local coordinates to global"""
    LocalMatrix=CreateLocalStiffnessMatrix(Beam)
    XN = Beam.StartNode.NodeCoor	[0]
    XK = Beam.EndNode.NodeCoor		[0]
    YN = Beam.StartNode.NodeCoor	[1]
    YK = Beam.EndNode.NodeCoor		[1]

    sina = (YK-YN)/math.sqrt((XK-XN)**2+(YK-YN)**2)
    cosa = (XK-XN)/math.sqrt((XK-XN)**2+(YK-YN)**2)
    TypeMatrix = {	1: numpy.array ([(cosa,  sina,    0 ),
                                           (-sina, cosa,    0 ),
                                           (  0,    0,      1 )
                                    ]),
                    2: numpy.array ([(cosa,  sina),
                                         ( -sina,  cosa)
                                    ])
                }

    #Pervaya matrica razmerom 3(2)x 5(6), so vtoroy soedinyaem po gorizontali

    XNM = TypeMatrix[Beam.StartNode.NodeType].shape [0]
    YNM = TypeMatrix[Beam.StartNode.NodeType].shape [1]
    XKM = TypeMatrix[Beam.EndNode.NodeType].shape   [0]
    YKM = TypeMatrix[Beam.EndNode.NodeType].shape   [1]

    A = numpy.array(TypeMatrix[Beam.StartNode.NodeType])
    B = numpy.array(TypeMatrix[Beam.EndNode.NodeType])
    A.resize((A.shape[0]+YKM, A.shape[1]), refcheck=False) 
    B = numpy.concatenate ((numpy.zeros((YNM, XKM)), B), 0)


    GradMatrix = numpy.concatenate((A, B), 1)
    #return GradMatrix


                                    # ])
    GlobalMatrix = GradMatrix*LocalMatrix
    #GlobalMatrix=numpy.dot(GradMatrix,LocalMatrix)
    return GlobalMatrix



node1=Nodes(1, 1, (1, 1))
node2=Nodes(2, 1, (2, 2))
beam1=Beams(1, node1, node2, 3, 2, 3)
matrix=CreateGlobalStiffnessMatrix(beam1)
if __name__=='__main__':
    print(matrix)
    print(matrix[1][1])